import torch
import torch.nn as nn
import torch.nn.functional as F
from vortex.utils.bbox import find_jaccard_overlap


def gcxgcy_to_cxcy(gcxgcy, priors_cxcy):
    xy = gcxgcy[:, :2] * priors_cxcy[:, 2:] / 10 + priors_cxcy[:, :2]
    wh = torch.exp(gcxgcy[:, 2:] / 5) * priors_cxcy[:, 2:]
    return torch.cat([xy, wh], dim=1)


def cxcy_to_xy(cxcy):
    xymin = cxcy[:, :2] - cxcy[:, 2:] / 2
    xymax = cxcy[:, :2] + cxcy[:, 2:] / 2
    return torch.cat([xymin, xymax], dim=1)


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class SSDDetect(object):
    def __init__(self, n_classes, priors, min_score, max_overlap, top_k):
        self.n_classes = n_classes
        self.priors = priors
        self.min_score = min_score
        self.max_overlap = max_overlap
        self.top_k = top_k
    
    def detect(self, predicted_locs, predicted_scores):
        """
        Decipher the 8732 locations ans class scores
        predicted_locs: [N, 8732, 4]
        predicted_scores: [N, 8732, ncls]
        """
        batch_size = predicted_locs.size(0)
        n_priors = self.priors.size(0)
        assert n_priors == predicted_locs.size(1)
        assert n_priors == predicted_scores.size(1)

        predicted_scores = F.softmax(predicted_scores, dim=2)

        all_images_boxes = []
        all_images_labels = []
        all_images_scores = []

        for i in range(batch_size):
            decoded_locs = cxcy_to_xy(gcxgcy_to_cxcy(predicted_locs[i], self.priors))
            # max_scores, best_label = predicted_scores[i].max(dim=1)

            image_boxes = []
            image_labels = []
            image_scores = []

            for c in range(1, self.n_classes):
                class_scores = predicted_scores[i][:, c]
                score_above_min_score = class_scores > self.min_score
                n_above_min_score = score_above_min_score.sum().item()
                if n_above_min_score == 0:
                    continue
                
                class_scores = class_scores[score_above_min_score]
                class_decoded_locs = decoded_locs[score_above_min_score]

                # sort predicted boxes and scores by score value
                class_scores, sort_ind = class_scores.sort(dim=0, descending=True)
                class_decoded_locs = class_decoded_locs[sort_ind]

                overlap = find_jaccard_overlap(class_decoded_locs, class_decoded_locs)
                
                # NMS
                suppress = torch.zeros((n_above_min_score), dtype=torch.uint8)
                for box in range(class_decoded_locs.size(0)):
                    if suppress[box] == 1:
                        continue
                    suppress = torch.max(suppress, overlap[box] > self.max_overlap)
                    suppress[box] = 0
                
                image_boxes.append(class_decoded_locs[1 - suppress])  # (N, 4)
                image_scores.append(class_scores[1 - suppress])  # (N, )
                image_labels.append(torch.LongTensor((1 - suppress).sum().item() * [c]))

            # concatenate into single tensor 
            if len(image_boxes) == 0:
                image_boxes.append(torch.FloatTensor([[0., 0., 1., 1.]]).to(device))
                image_labels.append(torch.LongTensor([0]).to(device))
                image_scores.append(torch.FloatTensor([0.]).to(device))

            image_boxes = torch.cat(image_boxes, dim=0)
            image_labels = torch.cat(image_labels, dim=0)
            image_scores = torch.cat(image_scores, dim=0)

            n_objects = image_scores.size(0)
            if n_objects > self.top_k:
                # keep only top_k objects
                image_scores, sort_ind = image_scores.sort(dim=0, descending=True)
                image_scores = image_scores[:top_k]
                image_boxes = image_boxes[sort_ind][:top_k]
                image_labels = image_labels[sort_ind][:top_k]
            all_images_boxes.append(image_boxes)
            all_images_labels.append(image_labels)
            all_images_scores.append(image_scores)
        
        return all_images_boxes, all_images_labels, all_images_scores
